// PNG image format loader
// (c) jimon game studio

#include "CImageFormatLoader_PNG.h"


#ifdef JE_SUPPORT_FILEFORMAT_PNG

#include "ILogManager.h"
#include "jeMemoryBuffer.h"
#include "png.h"
#include <sstream>
#include <ostream>
#include <iostream>

namespace je
{
	namespace video
	{
		static void png_cpexcept_error(png_structp png_ptr, png_const_charp msg)
		{
			JELOG_E(jeStringc("png fatal error : ")+jeStringc(msg))
			longjmp(png_ptr->jmpbuf, 1);
		}

		// PNG function for file reading
		void PNGAPI user_read_data_fcn(png_structp png_ptr, png_bytep data, png_size_t length)
		{
			core::IBuffer * DataBuffer =(core::IBuffer*)png_ptr->io_ptr;
			DataBuffer->Read(data,(u32)sizeof(c8)*length);
		}

		// PNG function for file writing
		void PNGAPI user_write_data_fcn(png_structp png_ptr, png_bytep data, png_size_t length)
		{
			std::stringstream * ststream = (std::stringstream*)png_ptr->io_ptr;
			ststream->write((const char*)data,(std::streamsize)length);
		}

		//! Constructor
		CImageFormatLoader_PNG::CImageFormatLoader_PNG()
		{
		}

		//! Destructor
		CImageFormatLoader_PNG::~CImageFormatLoader_PNG()
		{
		}

		//! Load Image
		IImage * CImageFormatLoader_PNG::LoadBuffer(core::IBuffer * DataBuffer)
		{
			u8** RowPointers = 0;

			DataBuffer->Lock(core::BLT_READ);

			u1 StreamWorkMode = DataBuffer->GetStreamWorkMode();
			DataBuffer->SetStreamWorkMode(true);
			DataBuffer->SetPosition(0);

			png_byte buffer[8];
			DataBuffer->Read(buffer,sizeof(png_byte)*8);

			if(png_sig_cmp(buffer, 0, 8))
			{
				JELOG_E("Can't load png file")
				DataBuffer->SetStreamWorkMode(StreamWorkMode);
				DataBuffer->UnLock();
				return NULL;
			}

			png_structp png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING,NULL, (png_error_ptr)png_cpexcept_error, NULL);
			if (!png_ptr)
			{
				JELOG_E("Internal PNG create read struct failure");
				DataBuffer->SetStreamWorkMode(StreamWorkMode);
				DataBuffer->UnLock();
				return NULL;
			}

			png_infop info_ptr = png_create_info_struct(png_ptr);
			if (!info_ptr)
			{
				JELOG_E("Internal PNG create info struct failure")
				png_destroy_read_struct(&png_ptr, NULL, NULL);
				DataBuffer->SetStreamWorkMode(StreamWorkMode);
				DataBuffer->UnLock();
				return NULL;
			}

			if (setjmp(png_jmpbuf(png_ptr)))
			{
				png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
				if (RowPointers)
					delete [] RowPointers;
				DataBuffer->SetStreamWorkMode(StreamWorkMode);
				DataBuffer->UnLock();
				return NULL;
			}

			png_set_read_fn(png_ptr, DataBuffer, user_read_data_fcn);

			png_set_sig_bytes(png_ptr, 8);

			png_read_info(png_ptr, info_ptr);

			u32 Width;
			u32 Height;
			s32 BitDepth;
			s32 ColorType;

			{
				png_uint_32 w,h;
				png_get_IHDR(png_ptr, info_ptr,&w, &h,&BitDepth, &ColorType, NULL, NULL, NULL);
				Width=w;
				Height=h;
			}

			if (ColorType==PNG_COLOR_TYPE_PALETTE)
				png_set_palette_to_rgb(png_ptr);

			if (BitDepth < 8)
			{
				if (ColorType==PNG_COLOR_TYPE_GRAY || ColorType==PNG_COLOR_TYPE_GRAY_ALPHA)
					png_set_gray_1_2_4_to_8(png_ptr);
				else
					png_set_packing(png_ptr);
			}

			if (png_get_valid(png_ptr, info_ptr, PNG_INFO_tRNS))
				png_set_tRNS_to_alpha(png_ptr);

			if (BitDepth == 16)
				png_set_strip_16(png_ptr);

			if (ColorType == PNG_COLOR_TYPE_GRAY || ColorType == PNG_COLOR_TYPE_GRAY_ALPHA)
				png_set_gray_to_rgb(png_ptr);

			if (ColorType == PNG_COLOR_TYPE_RGB)
				png_set_add_alpha(png_ptr,0xFF,PNG_FILLER_AFTER);

//			if (ColorType == PNG_COLOR_TYPE_RGB_ALPHA)
//				png_set_bgr(png_ptr);

			png_read_update_info(png_ptr, info_ptr);

			{
				png_uint_32 w,h;
				png_get_IHDR(png_ptr, info_ptr,&w, &h,&BitDepth, &ColorType, NULL, NULL, NULL);
				Width=w;
				Height=h;
			}

			//png_get_IHDR(png_ptr, info_ptr,(png_uint_32*)&Width, (png_uint_32*)&Height,&BitDepth, &ColorType, NULL, NULL, NULL);

			u8 BytesPerPixel = 4;

			u8 * DataArray = NULL;
			JENEW_A(DataArray,u8,Width*Height*BytesPerPixel)

			if (!DataArray)
			{
				JELOG_E("Internal PNG create image struct failure")
				png_destroy_read_struct(&png_ptr, NULL, NULL);
				DataBuffer->SetStreamWorkMode(StreamWorkMode);
				DataBuffer->UnLock();
				return NULL;
			}

			RowPointers = new png_bytep[Height];
			if (!RowPointers)
			{
				JELOG_E("Internal PNG create row pointers failure")
				png_destroy_read_struct(&png_ptr, NULL, NULL);
				JEDELETE_A(DataArray)
				DataBuffer->SetStreamWorkMode(StreamWorkMode);
				DataBuffer->UnLock();
				return NULL;
			}

			u8 * Data = DataArray;
			for (u32 i=0; i<Height; ++i)
			{
				RowPointers[i] = Data;
				Data += (Width*BytesPerPixel);
			}

			if (setjmp(png_jmpbuf(png_ptr)))
			{
				png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
				delete [] RowPointers;
				JEDELETE_A(DataArray)
				DataBuffer->SetStreamWorkMode(StreamWorkMode);
				DataBuffer->UnLock();
				return NULL;
			}

			png_read_image(png_ptr, RowPointers);

			png_read_end(png_ptr, NULL);
			delete [] RowPointers;
			png_destroy_read_struct(&png_ptr,&info_ptr, 0);

			DataBuffer->SetStreamWorkMode(StreamWorkMode);
			DataBuffer->UnLock();

			IImage::ImageColorFormat ColorFormat = IImage::ICF_R8G8B8A8;

			IImage * Image = CreateImage(jeVector2di(Width,Height),ColorFormat,DataArray);
			JEDELETE_A(DataArray)

			return Image;
		}

		//! Save Image
		void CImageFormatLoader_PNG::SaveBuffer(core::IBuffer * Buffer,IImage * Image)
		{
			if (!Buffer || !Image)
				return;

			std::stringstream outstream;

			png_structp png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING,NULL,(png_error_ptr)png_cpexcept_error,NULL);

			if (!png_ptr)
			{
				JELOG_E("Can't create png write struct")
				return;
			}

			png_infop info_ptr = png_create_info_struct(png_ptr);

			if (!info_ptr)
			{
				JELOG_E("Can't create png info struct")
				png_destroy_write_struct(&png_ptr, NULL);
				return;
			}

			if (setjmp(png_jmpbuf(png_ptr)))
			{
				JELOG_E("Can't set png jump")
				png_destroy_write_struct(&png_ptr, &info_ptr);
				return;
			}

			png_set_write_fn(png_ptr, &outstream, user_write_data_fcn, NULL);

			switch(Image->GetFormat())
			{
				case IImage::ICF_R8G8B8A8:
					png_set_IHDR(png_ptr,info_ptr,Image->GetSize().X,Image->GetSize().Y,8,PNG_COLOR_TYPE_RGB_ALPHA,PNG_INTERLACE_NONE,PNG_COMPRESSION_TYPE_DEFAULT,PNG_FILTER_TYPE_DEFAULT);
					break;
				default:
					png_set_IHDR(png_ptr,info_ptr,Image->GetSize().X,Image->GetSize().Y,8,PNG_COLOR_TYPE_RGB,PNG_INTERLACE_NONE,PNG_COMPRESSION_TYPE_DEFAULT,PNG_FILTER_TYPE_DEFAULT);
			}

			s32 lineWidth = Image->GetSize().X;
			switch(Image->GetFormat())
			{
				case IImage::ICF_R8G8B8A8:
					lineWidth *= 4;
					break;
			}

			u8 * tmpImage = NULL;
			JENEW_A(tmpImage,u8,Image->GetSize().Y * lineWidth)

			if (!tmpImage)
			{
				JELOG_E("Cant create png buff image")
				png_destroy_write_struct(&png_ptr, &info_ptr);
				return;
			}

			Image->Lock();

			switch(Image->GetFormat())
			{
				case IImage::ICF_R8G8B8A8:
					memcpy(tmpImage,Image->GetPtr(),Image->GetSize().Y * lineWidth);
					break;
			}

			Image->UnLock();

			u8 ** RowPointers = NULL;
			JENEW_A(RowPointers,png_bytep,Image->GetSize().Y)

			if (!RowPointers)
			{
				JELOG_E("Can't create png row pointers")
				png_destroy_write_struct(&png_ptr, &info_ptr);
				JEDELETE_A(tmpImage)
				return;
			}

			u8 * Data = tmpImage;

			for (s32 i=0; i< Image->GetSize().Y; ++i)
			{
				RowPointers[i] = Data;
				Data += lineWidth;
			}

			if (setjmp(png_jmpbuf(png_ptr)))
			{
				png_destroy_write_struct(&png_ptr, &info_ptr);
				JEDELETE_A(RowPointers)
				JEDELETE_A(tmpImage)
				return;
			}

			png_set_rows(png_ptr, info_ptr, RowPointers);

			if (Image->GetFormat() == IImage::ICF_R8G8B8A8)
				png_write_png(png_ptr, info_ptr, PNG_TRANSFORM_BGR, NULL);
			else
				png_write_png(png_ptr, info_ptr, PNG_TRANSFORM_IDENTITY, NULL);

			JEDELETE_A(RowPointers)
			JEDELETE_A(tmpImage)
			png_destroy_write_struct(&png_ptr, &info_ptr);

			s32 DataSize = (s32)outstream.str().size();
			std::string TempString = outstream.str();
			const c8 * WriteData = TempString.data();

			Buffer->Lock(core::BLT_WRITE);
			Buffer->Allocate(DataSize);
			Buffer->Write(WriteData,DataSize);
			Buffer->UnLock();

			return;
		}

		//! Return true if image loader can work with this data
		u1 CImageFormatLoader_PNG::CanWorkWithThis(core::IBuffer * DataBuffer)
		{
			if (!DataBuffer)
				return false;

			DataBuffer->Lock(core::BLT_READ);
			DataBuffer->SetPosition(0);

			png_byte buffer[8];
			DataBuffer->Read(buffer, 8);

			DataBuffer->UnLock();

			return !png_sig_cmp(buffer, 0, 8);
		}

		//! Return Image Format
		ImageFormat CImageFormatLoader_PNG::GetImageFormat()
		{
			return IF_PNG;
		}

	}
}

#endif
